/*
 * bpp_map_reduce_manyrhs.cc
 *
 * Copyright 2013 Long Tran (ltran3@gatech.edu)
 *
 * Block Principal Pivoting, solving in parallel for local minimum of
 *   L = sum |A - W H^T|^2
 *   subject to W >= 0, H >= 0
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#include <vector>
#include <armadillo>
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/mpi.hpp>

#include "helper/helper.h"
#include "helper/StopWatch.h"
#include "helper/FileAccess.h"
extern boost::mpi::communicator* p_world;
extern po::variables_map vm;
#include "helper/mpi_helper.h"
#include "helper/arma_helper.h"

using namespace arma;
#define MSG_TAG_CA 0
#define MSG_TAG_H 1

class BlockPrincipalPivotingMapReduceManyRHS
{
public:
  // input/output variables
  mat A, W, H;
  // iteration information
  std::vector<Iteration> iterInfo;
protected:
  // algorithm variables
  int m, n, k;
  double l1, l2;
  int maxiter, interval;
  bool normalizing;
  int bigiter;

  int start_col, end_col;
  std::vector<int> start_cols, end_cols;

  mat CC, CA, W_j, H_j, CC_i, CA_i;
  std::vector<mat> CA_split, H_split;

  StopWatch watch, clock;
  TimeDuration runtime[2], total_time[2];

public:
  BlockPrincipalPivotingMapReduceManyRHS() : CA_split(MPI_SIZE), H_split(MPI_SIZE) { 
    PO_REQUIRE(vm, "iter");
    PO_REQUIRE(vm, "k");

    PO_VARIABLE_VALUE(k, vm, "k", int);
    l1 = (vm.count("l1")) ? vm["l1"].as<double>() : 0; // l1 regularization
    l2 = (vm.count("l2")) ? vm["l2"].as<double>() : 0; // l2 regularization
    PO_VARIABLE_VALUE(maxiter, vm, "iter", int);
    PO_VARIABLE_VALUE(interval, vm, "interval", int);
    normalizing = vm.count("normalize");
  }

  void LoadData() 
  {
    PO_REQUIRE(vm, "data");
    PO_VARIABLE(data, vm, "data", std::string);
    
    std::stringstream ss;
    ss << data << "_" << MPI_RANK;
    A.load(ss.str(), arma_binary);
    m = A.n_rows;
    n = A.n_cols;
    //MSG("End Read data " << ss.str());
  }

  void Initialize()
  {
    W = randu<mat>(m, k);
    zero_threshold(W);

    if (MPI_RANK == 0) {
      H = randu<mat>(n, k);
      zero_threshold(H);
    }
    mpi::broadcast(*p_world, H, 0);
    
    //MSG("End random mat m=" << m << " n=" << n << " k=" << k << " MPI_SIZE=" << MPI_SIZE);
    // split similarly on rows of H on all nodes
    int split_n = n / MPI_SIZE;
    int start = 0;
    for (int i = 0; i < MPI_SIZE; i++, start += split_n) {
      start_cols.push_back(start);
      end_cols.push_back( (i < MPI_SIZE-1) ? start + split_n : n );
    }
    start_col = start_cols[MPI_RANK];
    end_col = end_cols[MPI_RANK];
    //MSG("End Initialize");
  }

  static void AbsThreshold(mat& v)
  {
    for (unsigned int i = 0; i < v.n_rows; i++)
      for (unsigned int j = 0; j < v.n_cols; j++)
	if (fabs(v(i,j)) < 1e-16) v(i,j) = 0;
  }

  static void Threshold(mat& v)
  {
    for (unsigned int i = 0; i < v.n_rows; i++)
      for (unsigned int j = 0; j < v.n_cols; j++)
	if (v(i,j) < 1e-16) v(i,j) = 0;
  }

  static void NegativeIndex(const mat& X, const umat& F, umat& H)
  {
    for (unsigned int i = 0; i < X.n_rows; i++)
      for (unsigned int j = 0; j < X.n_cols; j++)
	if ( F(i,j) && X(i, j) < -1e-16 ) H(i, j) = true;
  }

  void ExchangeIndex(unsigned int col, umat& F, const umat& H1, umat& G, const umat& H2)
  {
    for (unsigned int row = 0; row < F.n_rows; row++) {
      if (H1(row, col)) {
	F(row, col) = false;
	G(row, col) = true;
      }
      else if (H2(row, col)) {
	F(row, col) = true;
	G(row, col) = false;
      }
    }
  }

  void ExchangeMaxIndex(unsigned int col, umat& F, umat& G)
  {
    unsigned int row = F.n_rows-1;
    for (; row >= 0; row--)
      if ( F(row,col) || G(row,col) ) break;
    if (row >= 0) {
      F(row,col) = !F(row,col);
      G(row,col) = !G(row,col);
    }
  }

  struct ColumnCompare
  {
    const umat& m_F;
    ColumnCompare(const umat& F) : m_F(F) {}
    bool operator() (unsigned int i, unsigned int j) const {
      for (unsigned int row = 0; row < m_F.n_rows; row++)
	if (m_F(row,i) < m_F(row,j)) return true;
	else if (m_F(row,i) > m_F(row,j)) return false;
      return false;
    }
  };

  template<typename UVEC>
  static void GetMarker(const UVEC& marker, uvec& F)
  {
    int count = 0;
    for (unsigned int i=0; i < marker.n_elem; i++)
      if (marker(i)) count++;
    F.zeros(count);
    int j = 0;
    for (unsigned int i=0; i < marker.n_elem; i++)
      if (marker(i)) F(j++) = i;
  }

  // solve only on column in cols
  // F = index of non-active set (non zeros)
  void Solve(const uvec& cols, const umat& F, const mat& C, const mat& B, mat& X) 
  {
    //unsigned int k = F.n_rows;
    //unsigned int n = F.n_cols;
    // sort rows
    std::vector<unsigned int> col_idx(cols.n_elem);
    for (unsigned int i = 0; i < cols.n_elem; i++) col_idx[i] = cols(i);
    ColumnCompare comp(F);
    std::sort(ALL(col_idx), comp);

    // find break points and solve each chunk
    X.cols(cols).zeros();
    unsigned int start = 0;
    while (start < cols.n_elem) {
      unsigned int end = start;
      while ( end < cols.n_elem-1 && !comp(col_idx[end], col_idx[end+1]) ) end++;
      // from start to end are same passive set RHS
      uvec f_col = F.col( col_idx[start] ), f;
      GetMarker(f_col, f);
      uvec chunk_cols(end-start+1);
      for (unsigned int i = start; i <= end; i++) chunk_cols[i-start] = col_idx[i];

      mat chunk_X = solve( C.submat(f, f), B.submat(f, chunk_cols) );
      X.submat(f, chunk_cols) = chunk_X;

      start = end + 1;
    }
  }

  // C : k x k, B : k x n, X : k x n
  void SolveNNLSManyRHS(const mat& C, const mat& B, mat& X)
  {
    int pbar = 3;
    int q = C.n_cols;
    int n = B.n_cols;
    urowvec p = pbar*ones<urowvec>(n);
    urowvec t = (q + 1)*ones<urowvec>(n);
    
    umat F(q,n), G(q,n), H1(q,n), H2(q,n);
    F.zeros(); 
    G.ones(); // marker of indices in F and G
    X.zeros(q, n);  // initial value of X & Y
    mat Y = -B;

    //MSG(C.n_rows << " " << C.n_cols << " " << x.n_elem << " " << y.n_elem);
    double max_inner_iter = 5*q; // pow(2, q);
    for (int iter = 0; iter < max_inner_iter*pbar; iter++) {
      H1.zeros(); NegativeIndex(X, F, H1);
      H2.zeros(); NegativeIndex(Y, G, H2);

      //MSG("F=" << F.n_elem << " H1=" << H1.n_elem << " G=" << G.n_elem << " H2=" << H2.n_elem);
      urowvec h1ph2 = sum(H1) + sum(H2);
      uvec nonOpt;
      GetMarker(h1ph2, nonOpt);
      if (nonOpt.n_elem == 0) { // success
	//MSG("success");
	return;
      }
      else {
	for (int i = 0; i < n; i++) 
	  if ( h1ph2(i) > 0 && h1ph2(i) < t(i) ) {
	    t(i) = h1ph2(i);
	    p(i) = pbar;
	    ExchangeIndex(i, F, H1, G, H2);
	  }
	  else if ( h1ph2(i) > 0 && h1ph2(i) >= t(i) && p(i) >= 1 ) {
	    p(i) = p(i) - 1;
	    ExchangeIndex(i, F, H1, G, H2);
	  }
	  else if ( h1ph2(i) > 0 )   // backup rule
	    ExchangeMaxIndex(i, F, G);	
      }

      Solve(nonOpt, F, C, B, X);
      AbsThreshold(X);
      Y.cols(nonOpt) = C*X.cols(nonOpt)-B.cols(nonOpt);
      AbsThreshold(Y);
    }
    Threshold(X);
    MSG("failure");
  }

  void JoinH(const std::vector<mat>& H_split, mat& H)
  {
    for (int i = 0; i < MPI_SIZE; i++) {
      //MSG("Join " << start_cols[i] << " " <<  end_cols[i]-1 << " " << H.n_rows);
      H.rows(start_cols[i], end_cols[i]-1) = H_split[i];
    }
  }

  void SplitH(const arma::mat& H, std::vector<mat>& H_split)
  {
    for (int i = 0; i < MPI_SIZE; i++) {
      //MSG("Split " << start_cols[i] << " " <<  end_cols[i]-1 << " " << H.n_cols);
      H_split[i] = trans(H.cols(start_cols[i], end_cols[i]-1));
    }
  }

  void UpdateW()
  {
    // update W
    CC = trans(H)*H;        // local, same in all nodes
    CA = trans(H)*trans(A); // size k x m_i
    W_j = zeros<mat>(k, m);
    SolveNNLSManyRHS(CC, CA, W_j);
    W = trans(W_j);
  }

  void SendCAFromMaster()
  {
    if (MPI_RANK == 0) { // send part of CA to slave nodes
      //MSG("CA.m=" << CA.n_rows << " CA.n=" << CA.n_cols);
      mpi::request reqs[MPI_SIZE-1];
      SplitH(CA, CA_split);
      for (int j = 1; j < MPI_SIZE; j++)
	MPI_ISEND_VERBOSE(reqs[j-1], j, MSG_TAG_CA, CA_split[j], vm, "CA["<<j<<"]");
      mpi::wait_all(reqs, reqs+MPI_SIZE-1);
      CA = trans(CA_split[0]);
    }
    else {
      mpi::request reqs;
      MPI_IRECV_VERBOSE(reqs, 0, MSG_TAG_CA, CA, vm, "CA["<<MPI_RANK<<"]");
      reqs.wait();
      CA = trans(CA);
    }    
  }

  void ReceiveHjFromSlaves()
  {
    if (MPI_RANK == 0) {
      mpi::request reqs[MPI_SIZE-1];
      H_split[0] = H_j;
      for (int j = 1; j < MPI_SIZE; j++)
	MPI_IRECV_VERBOSE(reqs[j-1], j, MSG_TAG_H, H_split[j], vm, "H["<<j<<"]");
      mpi::wait_all(reqs, reqs+MPI_SIZE-1);
      JoinH(H_split, H);
    }
    else {
      mpi::request reqs;
      MPI_ISEND_VERBOSE(reqs, 0, MSG_TAG_H, H_j, vm, "H["<<MPI_RANK<<"]");
      reqs.wait();
    }
    mpi::broadcast(*p_world, H, 0);
  }

  void UpdateH(StopWatch& clock, TimeDuration runtime[])
  {
    clock.GetDuration();
    CC_i = trans(W)*W; // k x k
    CA_i = trans(W)*A; // k x n
    runtime[0] += clock.GetDuration();
    
    clock.GetDuration();
    mpi::all_reduce(*p_world, CC_i, CC, std::plus<mat>());
    mpi::reduce(*p_world, CA_i, CA, std::plus<mat>(), 0);
    SendCAFromMaster();
    //MSG(MPI_RANK << ": CA.m=" << CA.n_rows << " CA.n=" << CA.n_cols);
    runtime[1] += clock.GetDuration();
    
    clock.GetDuration();
    H_j = zeros<mat>(k, end_col - start_col);
    SolveNNLSManyRHS(CC, CA, H_j);
    H_j = trans(H_j);
    runtime[0] += clock.GetDuration();
    
    // concatenate parts of H from H_j in master node then broadcast
    clock.GetDuration();
    ReceiveHjFromSlaves();
    runtime[1] += clock.GetDuration();
  }

  void RecordProgress(int bigiter)
  {
    double res = norm(A-W*trans(H), "fro"), sum_res;
    res *= res;
    mpi::reduce(*p_world, res, sum_res, std::plus<double>(), 0);
    
    if (MPI_RANK == 0) {
      Iteration info(bigiter, sqrt(sum_res), runtime[0]+runtime[1], runtime[0], runtime[1]);
      iterInfo.push_back(info);
      if (bigiter % interval == 0) {
	MSG("it=" << bigiter << " |res|=" << sqrt(sum_res)  << " |H|=" << norm(H, "fro")
	    << " d=" << SECONDS(info.duration));
      }
    }
  }

  void SolveBPP()
  {
    Initialize();
    
    watch.Reset();
    clock.Reset();
    for (bigiter = 0; bigiter < maxiter; bigiter++) {
      RecordProgress(bigiter);

      clock.GetDuration();
      UpdateW();
      runtime[0] += clock.GetDuration();

      UpdateH(clock, runtime);
      
      // print information
    }
    RecordProgress(maxiter);
    if (MPI_RANK == 0) MSG(watch.GetTotalDuration());
    mpi::reduce(*p_world, runtime[0], total_time[0], std::plus<TimeDuration>(), 0);
    mpi::reduce(*p_world, runtime[1], total_time[1], std::plus<TimeDuration>(), 0);
    if (MPI_RANK == 0)
      for (int i = 0; i < 2; i++) MSG("clock " << i << " " << total_time[i]);
  }
};

void run_map_reduce_bpp_manyrhs()
{
  if (MPI_RANK == 0) MSG("Master/Slave Block Principal Pivoting NMF with many RHS");
  BlockPrincipalPivotingMapReduceManyRHS a;
  a.LoadData();
  a.SolveBPP();
  
  if (MPI_RANK == 0 && vm.count("output")) {
    PO_VARIABLE(output, vm, "output", std::string);
    a.H.save(output, raw_ascii);
  }

  if (MPI_RANK == 0 && vm.count("iterinfo")) {
    PO_VARIABLE(infoFile, vm, "iterinfo", std::string);
    FileAccessWrite(out, infoFile, NULL);
    const std::vector<Iteration>& info =  a.iterInfo;
    for (unsigned int i = 0; i < info.size(); i++) {
      out << info[i].iter << "\t" << (double)info[i].duration.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].compute.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].communication.total_milliseconds()/1000 << "\t"
	  << info[i].residual << std::endl;
    }
  }
}
